import re
import json
from datetime import datetime
from llama_index.core.llms import ChatMessage,MessageRole
from biz.data.sql_template.sql_template_store import SqlTemplateStore

class SQLGenerator:
    def __init__(self, llm):
        self.llm = llm

    def generate_sql(self, template, question):
        """
        获取模板并生成SQL

        Args:
            template: SQL模板
            question: 用户问题

        Returns:
            生成的SQL查询
        """
        sql_template_prompt = """
        你是一个智能问题匹配系统，负责从候选问题集合中选出与用户提出的问题语义最接近的候选项，并根据匹配结果生成相关输出。如果未找到符合匹配阈值的问题，返回格式化的空结果。请严格按照以下要求完成任务：

        1. **输入参数**：
           - `user_question`：一个字符串，表示用户提出的问题。
           - `candidate_questions`：一个 JSON 数组，其中每个对象包含以下字段：
             - `id`（字符串）：候选问题的唯一标识。
             - `description`（字符串）：候选问题的文本内容。
             - `sql_template`：一个包含占位符的字符串，用于生成最终的 SQL 查询语句。
        
        2. **执行逻辑**：
           - 对 `user_question` 与 `candidate_questions` 中每个问题的 `description` 进行语义相似度计算。
           - 找出语义相似度最高的候选问题，并记录其相似度得分。
           - 如果最高相似度得分小于 0.8（匹配阈值），则判定为无合理匹配，返回格式化的空结果。
           - 如果找到匹配问题，使用匹配结果中的相关参数更新 `sql_template`，生成完整的 SQL 查询语句。
           - 对生成的 SQL 语句进行语法检查。
        
        3. **输出格式**：
           - 返回结果为 JSON 对象，字段顺序和格式如下：
             ```json
             {
               "matched_question_id": "最匹配的候选问题id（字符串）",
               "match_score": 匹配度得分（0 到 1 的浮点数，越高表示越相似）,
               "is_matched": true 或 false（布尔值，表示是否找到合理匹配）,
               "parameters": 相关参数的 JSON 对象（如无匹配则为空对象）,
               "sql": 替换参数后生成的完整 SQL 查询语句（保证语法正确，如未找到匹配则为空字符串）
             }
             ```
        
        4. **输出规则**：
           - 当 `is_matched` 为 `false` 时：
             - `matched_question_id` 必须为 `""`（空字符串）。
             - `match_score` 必须为 `0`。
             - `parameters` 必须为 `{}`（空对象）。
             - `sql` 必须为 `""`（空字符串）。
           - SQL中可选条件不需要时，不要输出。
           - 输出必须符合 JSON 规范，严格按照字段定义及顺序生成，不得遗漏字段或添加多余字段。
        
        5. **注意事项**：
           - 语义相似度计算方法需高效且支持多语言对比。
           - 生成的 SQL 查询语句需完整替换所有参数并确保符合 PostgreSQL 语法规范。
           - 语法检查和修正功能需严谨，修正后的 SQL 必须符合规范且不能为空。
           - 输出结果需保持一致性和完整性，对于失败或无匹配的情况，遵循定义的空结果格式返回。
           - 不要进行解释或说明，直接返回 JSON 格式的结果。
        """

        sql_template_user_prompt = f"""
        输入信息：
        - current_date：{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
        - user_question："{question}"
        - candidate_questions：
            {template}
        """
        messages = [
            ChatMessage(
                role=MessageRole.SYSTEM,
                content=sql_template_prompt,
            ),
            ChatMessage(
                role=MessageRole.USER,
                content=sql_template_user_prompt,
            ),
        ]
        try:
            response = self.llm.chat(messages)
            print(f"SQLGenerator: {response.message.content}")
            content = response.message.content
            data = json.loads(self.extract_json_from_response(content))
            return data
        except Exception as e:
            print(f"Error generate sql: {e}")
            return None

    def extract_json_from_response(self, content):
        """
        从LLM响应中提取JSON内容
        """
        try:
            # 方法1: 提取```json```代码块
            json_block_pattern = r'```json\s*(\{.*?\})\s*```'
            json_block_match = re.search(json_block_pattern, content, re.DOTALL)
            if json_block_match:
                return json_block_match.group(1).strip()

            # 方法2: 查找独立的JSON对象（使用括号计数确保完整性）
            start_idx = content.find('{')
            if start_idx == -1:
                return None

            brace_count = 0
            for i, char in enumerate(content[start_idx:], start_idx):
                if char == '{':
                    brace_count += 1
                elif char == '}':
                    brace_count -= 1
                    if brace_count == 0:
                        return content[start_idx:i+1]

            return None
        except Exception as e:
            print(f"JSON提取错误: {e}")
            return None

    def generate_plotly_code(self, template, id=None):
        """
        使用提取的参数填充Plotly代码模板

        Args:
            template: 包含Plotly代码的模板
        Returns:
            生成的Plotly代码，如果模板中没有则返回None
        """
        try:
            if not id:
                return None

            # 解析模板（如果是字符串格式的JSON）
            if isinstance(template, str):
                template_data = json.loads(template)
            else:
                template_data = template

            for item in template_data:
                if item.get('id') == id:
                    return item.get('plotly_code')
            return None
        except Exception as e:
            print(f"Error generating plotly code: {e}")
            return None
